import os
from typing import List
import numpy as np
import cv2
from PIL import Image, ImageDraw
from tqdm import tqdm
from xml.dom.minidom import parse

dict_class = {'ship': 0, 'plane': 1}
dict_color = {'ship': (255, 0, 0), 'plane': (0, 0, 255)}


def get_xml_data(xml_path, img_path, annotated_img_path, save_txt_path, size=1024):
    dom = parse(xml_path)
    root = dom.documentElement
    # img_name = root.getElementsByTagName("filename")[0].childNodes[0].data
    # img_size = root.getElementsByTagName("size")[0]
    objects = root.getElementsByTagName("object")
    # img_w = float(img_size.getElementsByTagName("width")[0].childNodes[0].data)
    # img_h = float(img_size.getElementsByTagName("height")[0].childNodes[0].data)
    # img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
    # print("img_name:", img_name)
    # print("image_info:(w,h,c)", img_w, img_h, img_c)
    f = open(save_txt_path, 'w')
    img = Image.open(img_path)
    draw = ImageDraw.Draw(img)
    for box in objects:
        cls_name = box.getElementsByTagName("name")[0].childNodes[0].data
        x0 = float(box.getElementsByTagName("x0")[0].childNodes[0].data)
        y0 = float(box.getElementsByTagName("y0")[0].childNodes[0].data)
        x1 = float(box.getElementsByTagName("x1")[0].childNodes[0].data)
        y1 = float(box.getElementsByTagName("y1")[0].childNodes[0].data)
        x2 = float(box.getElementsByTagName("x2")[0].childNodes[0].data)
        y2 = float(box.getElementsByTagName("y2")[0].childNodes[0].data)
        x3 = float(box.getElementsByTagName("x3")[0].childNodes[0].data)
        y3 = float(box.getElementsByTagName("y3")[0].childNodes[0].data)
        draw.line([(x0, y0), (x1, y1), (x2, y2), (x3, y3), (x0, y0)], fill=dict_color.get(cls_name), width=3)
        line = trans([dict_class.get(cls_name, 0), x0, y0, x1, y1, x2, y2, x3, y3], size)
        if line:
            f.write(line)
    f.close()
    img.save(annotated_img_path)


def trans(box: List, size=1024):
    cls = box[0]
    data = np.array(box[1:])
    data = data.reshape(4, 2).astype(int)  # 这里必须是整数
    rect = cv2.minAreaRect(data)  # 得到最小外接矩形的（中心(x,y), (宽,高), 旋转角度）
    c_x = rect[0][0]
    c_y = rect[0][1]
    w = rect[1][0]
    h = rect[1][1]
    new_box = np.array([c_x, c_y, w, h])
    if (new_box < 0).any() or (new_box > size).any():  # 坐标为负数，说明在边界处，或者box宽度超过图片尺寸，直接跳过
        return None
    theta = int(rect[-1])
    long_side, short_side = w, h
    if w < h:
        long_side = h
        short_side = w
        theta -= 90
    theta += 90
    line = f'{cls} {c_x / size} {c_y / size} {long_side / size} {short_side / size} {theta}\n'
    return line


def transfer_xmls(images_dir, annotated_images_dir, labels_dir, save_dir, size=1024):
    # os.makedirs(save_dir, exist_ok=True)
    os.makedirs(annotated_images_dir, exist_ok=True)
    for file_name in tqdm(os.listdir(images_dir)):
        xml_path = os.path.join(labels_dir, file_name.replace('.png', '.xml'))
        img_path = os.path.join(images_dir, file_name)
        annotated_img_path = os.path.join(annotated_images_dir, file_name.replace('.png', '_annotated.png'))
        # save_txt_path = os.path.join(save_dir, file_name.replace('.png', '.txt'))
        save_txt_path = ''
        get_xml_data(xml_path, img_path, annotated_img_path, save_txt_path, size)


def remove_none_txt(label_txt_dir, image_dir):
    for file_name in tqdm(os.listdir(label_txt_dir)):
        txt_path = os.path.join(label_txt_dir, file_name)
        with open(txt_path) as f:
            data = f.read()
        if not data:
            os.remove(txt_path)
            print(f'成功删除：{txt_path}')
            image_path = os.path.join(image_dir, f'{file_name.strip(".txt")}.png')
            os.remove(image_path)
            print(f'成功删除：{image_path}')


def get_txt_data(txt_path, img_path, new_img_path, annotated_img_path, save_txt_path, size):
    """
    将类似 x1 y1 x2 y2 x3 y3 x4 y4 classname id 这种格式的txt标注文件转换成yolov5旋转框需要的格式:
    id center_x center_y long_side short_side angle
    坐标进行归一化处理
    将原始训练图片统一大小，并生成标注图片
    :param txt_path:
    :param img_path:
    :param annotated_img_path:
    :param save_txt_path:
    :param size:
    :return:
    """
    img_array = np.array(Image.open(img_path))
    img_resized_array, scale, new_w, new_h = uniform_img_size(img_array, size)
    img_resized = Image.fromarray(img_resized_array)
    img_resized.save(new_img_path)
    draw = ImageDraw.Draw(img_resized)
    f = open(txt_path)
    f2 = open(save_txt_path, 'w')
    data = f.readlines()
    for line in data:
        box = line.strip().split()
        cls_name = box[-2]
        cls_id = box[-1]
        box = list(map(float, box[:8]))
        xs = np.array(box[::2]) * scale + (size - new_w) // 2
        ys = np.array(box[1::2]) * scale + (size - new_h) // 2
        draw.line([(xs[0], ys[0]), (xs[1], ys[1]), (xs[2], ys[2]), (xs[3], ys[3]), (xs[0], ys[0])],
                  fill=dict_color.get(cls_name), width=3)
        line = trans([cls_id, xs[0], ys[0], xs[1], ys[1], xs[2], ys[2], xs[3], ys[3]], size)
        if line:
            f2.write(line)
    f.close()
    f2.close()
    img_resized.save(annotated_img_path)


def uniform_img_size(img_array, size):
    """
    将训练图片统一尺寸
    保持长宽比不变，周围填充0以满足输出的大小要求
    :param img_array:
    :param size: 图片大小，宽高一样，1024
    :return:
    """
    img_w, img_h = img_array.shape[1], img_array.shape[0]
    scale = size / max(img_h, img_w)
    new_w = int(img_w * scale)
    new_h = int(img_h * scale)
    resized_image = cv2.resize(img_array, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    canvas = np.full((size, size, 3), 0).astype(np.uint8)
    canvas[(size - new_h) // 2:(size - new_h) // 2 + new_h, (size - new_w) // 2:(size - new_w) // 2 + new_w,
    :] = resized_image
    return canvas, scale, new_w, new_h


def transfer_labeltxt(images_dir, new_img_dir, annotated_images_dir, labels_dir, save_dir, img_suffix='.bmp',
                      size=1024):
    os.makedirs(save_dir, exist_ok=True)
    os.makedirs(annotated_images_dir, exist_ok=True)
    os.makedirs(new_img_dir,exist_ok=True)
    for file_name in tqdm(os.listdir(images_dir)):
        txt_path = os.path.join(labels_dir, file_name.replace(img_suffix, '.txt'))
        img_path = os.path.join(images_dir, file_name)
        new_img_path = os.path.join(new_img_dir, file_name)
        annotated_img_path = os.path.join(annotated_images_dir,
                                          file_name.replace(img_suffix, f'_annotated{img_suffix}'))
        save_txt_path = os.path.join(save_dir, file_name.replace(img_suffix, '.txt'))
        get_txt_data(txt_path, img_path, new_img_path, annotated_img_path, save_txt_path, size)


if __name__ == '__main__':
    img_dir = r'F:\dataset\object_detection\dota_data\Alldotaimg'
    new_img_dir = r'F:\dataset\object_detection\dota_data\images'
    annotated_images_dir = r'F:\dataset\object_detection\dota_data\img_annotated'
    txt_dir = r'F:\dataset\object_detection\dota_data\Alldoatlabel'
    save_dir = r'F:\dataset\object_detection\dota_data/labels'
    transfer_labeltxt(img_dir, new_img_dir, annotated_images_dir, txt_dir, save_dir)
    # transfer_xmls(img_dir, annotated_images_dir, xml_dir, save_dir)
# remove_none_txt(label_txt_dir=r'E:\workspace\rotation-yolov5\dataset\dota1.5\rotated_data\train_data\labels\val',
#                 image_dir=r'E:\workspace\rotation-yolov5\dataset\dota1.5\rotated_data\train_data\images\val')
